// Copyright (c) Aptos Foundation
// Licensed pursuant to the Innovation-Enabling Source Code License, available at https://github.com/aptos-labs/aptos-core/blob/main/LICENSE

use crate::{
    counters::{self, MAX_TXNS_FROM_BLOCK_TO_EXECUTE, TXNS_IN_BLOCK, TXN_SHUFFLE_SECONDS},
    payload_manager::TPayloadManager,
    transaction_deduper::TransactionDeduper,
    transaction_shuffler::TransactionShuffler,
};
use aptos_config::config::BlockTransactionFilterConfig;
use aptos_consensus_types::{block::Block, quorum_cert::QuorumCert};
use aptos_crypto::HashValue;
use aptos_executor_types::ExecutorResult;
use aptos_types::transaction::SignedTransaction;
use fail::fail_point;
use futures::future::Shared;
use move_core_types::account_address::AccountAddress;
use std::{future::Future, sync::Arc, time::Instant};

pub struct BlockPreparer {
    payload_manager: Arc<dyn TPayloadManager>,
    txn_filter_config: Arc<BlockTransactionFilterConfig>,
    txn_deduper: Arc<dyn TransactionDeduper>,
    txn_shuffler: Arc<dyn TransactionShuffler>,
}

impl BlockPreparer {
    pub fn new(
        payload_manager: Arc<dyn TPayloadManager>,
        txn_filter_config: Arc<BlockTransactionFilterConfig>,
        txn_deduper: Arc<dyn TransactionDeduper>,
        txn_shuffler: Arc<dyn TransactionShuffler>,
    ) -> Self {
        Self {
            payload_manager,
            txn_filter_config,
            txn_deduper,
            txn_shuffler,
        }
    }

    pub async fn prepare_block(
        &self,
        block: &Block,
        block_qc_fut: Shared<impl Future<Output = Option<Arc<QuorumCert>>>>,
    ) -> ExecutorResult<(Vec<SignedTransaction>, Option<u64>)> {
        fail_point!("consensus::prepare_block", |_| {
            use aptos_executor_types::ExecutorError;
            use std::{thread, time::Duration};
            thread::sleep(Duration::from_millis(10));
            Err(ExecutorError::CouldNotGetData)
        });
        let start_time = Instant::now();

        let (txns, max_txns_from_block_to_execute, block_gas_limit) = tokio::select! {
                // Poll the block qc future until a QC is received. Ignore None outcomes.
                Some(qc) = block_qc_fut => {
                    let block_voters = Some(qc.ledger_info().get_voters_bitvec().clone());
                    self.payload_manager.get_transactions(block, block_voters).await
                },
                result = self.payload_manager.get_transactions(block, None) => {
                   result
                }
        }?;
        TXNS_IN_BLOCK
            .with_label_values(&["before_filter"])
            .observe(txns.len() as f64);

        let txn_filter_config = self.txn_filter_config.clone();
        let txn_deduper = self.txn_deduper.clone();
        let txn_shuffler = self.txn_shuffler.clone();

        let block_id = block.id();
        let block_author = block.author();
        let block_epoch = block.epoch();
        let block_timestamp_usecs = block.timestamp_usecs();

        // Transaction filtering, deduplication and shuffling are CPU intensive tasks, so we run them in a blocking task.
        let result = tokio::task::spawn_blocking(move || {
            let filtered_txns = filter_block_transactions(
                txn_filter_config,
                block_id,
                block_author,
                block_epoch,
                block_timestamp_usecs,
                txns,
            );
            let deduped_txns = txn_deduper.dedup(filtered_txns);
            let mut shuffled_txns = {
                let _timer = TXN_SHUFFLE_SECONDS.start_timer();

                txn_shuffler.shuffle(deduped_txns)
            };

            if let Some(max_txns_from_block_to_execute) = max_txns_from_block_to_execute {
                shuffled_txns.truncate(max_txns_from_block_to_execute as usize);
            }
            TXNS_IN_BLOCK
                .with_label_values(&["after_filter"])
                .observe(shuffled_txns.len() as f64);
            MAX_TXNS_FROM_BLOCK_TO_EXECUTE.observe(shuffled_txns.len() as f64);
            Ok(shuffled_txns)
        })
        .await
        .expect("Failed to spawn blocking task for transaction generation");
        counters::BLOCK_PREPARER_LATENCY.observe_duration(start_time.elapsed());
        result.map(|result| (result, block_gas_limit))
    }
}

/// Filters transactions in a block based on the filter configuration
fn filter_block_transactions(
    txn_filter_config: Arc<BlockTransactionFilterConfig>,
    block_id: HashValue,
    block_author: Option<AccountAddress>,
    block_epoch: u64,
    block_timestamp_usecs: u64,
    txns: Vec<SignedTransaction>,
) -> Vec<SignedTransaction> {
    // If the transaction filter is disabled, return early
    if !txn_filter_config.is_enabled() {
        return txns;
    }

    // Otherwise, filter the transactions
    txn_filter_config
        .block_transaction_filter()
        .filter_block_transactions(
            block_id,
            block_author,
            block_epoch,
            block_timestamp_usecs,
            txns,
        )
}
